#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import opengate as gate
import linac_elekta_synergy
import uproot
from pathlib import Path

if __name__ == "__main__":
    # create the simulation
    sim = gate.Simulation()

    # main options
    sim.g4_verbose = False
    # sim.visu = True
    sim.visu_type = "vrml"
    sim.number_of_threads = 2
    sim.random_seed = "auto"
    # sim.random_seed = 123456
    sim.progress_bar = True
    sim.output_dir = "./output"

    # units
    m = gate.g4_units.m
    cm = gate.g4_units.cm
    mm = gate.g4_units.mm
    nm = gate.g4_units.nm
    Bq = gate.g4_units.Bq
    MeV = gate.g4_units.MeV

    # world size
    sim.world.size = [1 * m, 1 * m, 3 * m]
    sim.world.material = "G4_AIR"

    # waterbox
    wb = sim.add_volume("Box", "waterbox")
    wb.size = [30 * cm, 30 * cm, 30 * cm]
    wb.material = "G4_WATER"
    wb.color = [0, 0, 1, 1]  # blue

    # virtual plane for phase space (no linac !)
    plane = sim.add_volume("Tubs", "phase_space_plane")
    plane.material = "G4_AIR"
    plane.rmin = 0
    plane.rmax = 100 * mm
    plane.dz = 1 * nm
    # the position is relative to the 'linac' volume
    # that was 100 cm away.
    plane.translation = [0, 0, -27 * cm + 100 * cm]
    plane.color = [1, 0, 0, 1]  # red

    # add MLC
    mlc = linac_elekta_synergy.add_MLC(sim, sim.world.name, "mlc", [7 * cm, 7 * cm])
    mlc.translation[2] += 100 * cm

    # consider the "large" phase space (1M events, about 296 MB)
    data_folder = Path("data")
    phsp_filename = "phsp_1M.root"

    # phsp source
    source = sim.add_source("PhaseSpaceSource", "phsp_source")
    source.attached_to = plane.name
    source.phsp_file = data_folder / phsp_filename
    source.position_key = "PrePositionLocal"
    source.direction_key = "PreDirectionLocal"
    source.energy_key = "KineticEnergy"
    source.weight_key = "Weight"
    source.global_flag = False
    source.particle = "gamma"
    source.n = 2000000 / sim.number_of_threads
    source.batch_size = source.n
    if sim.number_of_threads > 1:
        source.entry_start = [e * source.n for e in range(sim.number_of_threads)]
        phsp = uproot.open(source.phsp_file)
        first_branch = phsp.keys()[0]
        phsp_n = int(phsp[first_branch].num_entries)
        print(f"The phsp file {source.phsp_file} contains {phsp_n} values")
        print(f"The threads will start in the phsp at position {source.entry_start}")
    print(f"Each thread will simulate {source.n} particles")

    # add stat actor
    stats = sim.add_actor("SimulationStatisticsActor", "stats")
    stats.track_types_flag = True
    stats.output_filename = "stats3.txt"

    # Dose actor (edep only)
    dose = sim.add_actor("DoseActor", "dose")
    dose.output_filename = "waterbox2.mhd"
    dose.attached_to = "waterbox"
    dose.size = [30, 30, 30]
    dose.spacing = [10 * mm, 10 * mm, 10 * mm]
    dose.edep_uncertainty.active = True
    dose.hit_type = "random"

    # phys
    sim.physics_manager.physics_list_name = "G4EmStandardPhysics_option3"
    sim.physics_manager.set_production_cut("world", "all", 1 * mm)

    # ---------------------------------------------------------------------
    # start simulation
    sim.run()

    # end
    print(stats)
